#!/usr/bin/perl -w
# Author: Xinwu
# Describe: 
# Completion Time: 
# Email: lexinwu@outlook.com
use strict;
use Getopt::Long;
use File::Path qw/mkpath/;

sub USAGE {
    my $usage=<<"USAGE";
Usage: perl $0 --mtx <matrix.csv> --outdir <outdir> --tax <human or mouse> [--db_dir db_dir] [--num_workers num_workers] [--transpose ] 
Options:
    --mtx                   expression matrix for the single cell experiment, formats are supported: csv(rows=cells x columns=genes).
    --transpose             transpose the expression matrix (rows=genes x columns=cells).
    --outdir                result output dir.
    --tax                   the conventional species such as human,mouse.
    --num_workers           the number of workers to use(default: 16).
    --db_dir                where tf_file feather_file and tbl_file saved.
    -h                      print help messages.

Database directory structure
    ├── hg19-500bp-upstream-10species.mc9nr.genes_vs_motifs.rankings.feather
    ├── hg19-tss-centered-10kb-10species.mc9nr.genes_vs_motifs.rankings.feather
    ├── hs_hgnc_tfs.txt
    ├── mm9-500bp-upstream-10species.mc9nr.genes_vs_motifs.rankings.feather
    ├── mm9-tss-centered-10kb-10species.mc9nr.genes_vs_motifs.rankings.feather
    ├── mm_mgi_tfs.txt
    ├── motifs-v9-nr.hgnc-m0.001-o0.0.tbl
    └── motifs-v9-nr.mgi-m0.001-o0.0.tbl

USAGE
        print $usage;
        exit;
}


my ($mtx, $transpose, $outdir, $tax, $num_workers, $db_dir);
GetOptions(
    "h|help"                => \&USAGE, 
    "mtx:s"                 => \$mtx, 
    "transpose!"            => \$transpose,
    "outdir:s"              => \$outdir,
    "tax:s"                 => \$tax,
    "num_workers:i"         => \$num_workers,
    "db_dir:s"              => \$db_dir,
) or &USAGE;
&USAGE unless ($mtx and $outdir and $tax);

sub format_datetime {
    my($sec, $min, $hour, $day, $mon, $year, $wday, $yday, $isdst) = @_;
    $wday = $yday = $isdst = 0;
    sprintf("%4d-%02d-%02d %02d:%02d:%02d", $year+1900, $mon+1, $day, $hour, $min, $sec);
}


$db_dir         ||= "/public/home/lxw/10.Git/SCENIC/data";
$num_workers    ||= 16;
my $pyscenic = `which pyscenic`;
chomp $pyscenic;  # 遇到个怪事 which pyscenic 的时候, 它自己加了个回车, 导致命令不在一行而报错
chop $outdir if $outdir =~ /\/$/;
chop $db_dir if $db_dir =~ /\/$/;
mkpath $outdir unless -e $outdir;

my($tf,$upstream,$tss,$tbl);
if($tax eq "mouse"){
    $tf         = "$db_dir/mm_mgi_tfs.txt";
    $upstream   = "$db_dir/mm9-500bp-upstream-10species.mc9nr.genes_vs_motifs.rankings.feather";
    $tss        = "$db_dir/mm9-tss-centered-10kb-10species.mc9nr.genes_vs_motifs.rankings.feather";
    $tbl        = "$db_dir/motifs-v9-nr.mgi-m0.001-o0.0.tbl";
}elsif($tax eq "human"){
    $tf         = "$db_dir/hs_hgnc_tfs.txt";
    $upstream   = "$db_dir/hg19-500bp-upstream-10species.mc9nr.genes_vs_motifs.rankings.feather";
    $tss        = "$db_dir/hg19-tss-centered-10kb-10species.mc9nr.genes_vs_motifs.rankings.feather";
    $tbl        = "$db_dir/motifs-v9-nr.hgnc-m0.001-o0.0.tbl";
}else{
    print "Currently, only human and mouse are supported, not other species\n";
    exit;
}


if($transpose){
    my $time = &format_datetime(localtime(time()));
    my $CMD1 = "$pyscenic grn --method grnboost2 --output $outdir/step1_mx_adjacencies.tsv --num_workers $num_workers --seed 100 --transpose $mtx $tf";
    print "[INFO: $time] Now start run step1 may take a long time, please be patient. \n";
    print "CMD: $CMD1\n";
    system("$CMD1 >> $outdir/run_log.txt 2>\&1");
}else{
    my $time = &format_datetime(localtime(time()));
    my $CMD1 = "$pyscenic grn --method grnboost2 --output $outdir/step1_mx_adjacencies.tsv --num_workers $num_workers --seed 100 $mtx $tf";
    print "[INFO: $time] Now start run step1 may take a long time, please be patient. \n";
    print "CMD: $CMD1\n";
    system("$CMD1 >> $outdir/run_log.txt 2>\&1");
}

if($transpose){
    my $time = &format_datetime(localtime(time()));
    my $CMD2 = "$pyscenic ctx $outdir/step1_mx_adjacencies.tsv $upstream $tss --annotations_fname $tbl --expression_mtx_fname $mtx --mode dask_multiprocessing" .
                " --output $outdir/step2_regulons.csv --num_workers $num_workers --transpose";
    print "[INFO: $time] Now start run step2 may take a long time, please be patient. \n";
    print "CMD: $CMD2\n";
    system("$CMD2 >> $outdir/run_log.txt 2>\&1");
}else{
    my $time = &format_datetime(localtime(time()));
    my $CMD2 = "$pyscenic ctx $outdir/step1_mx_adjacencies.tsv $upstream $tss --annotations_fname $tbl --expression_mtx_fname $mtx --mode dask_multiprocessing" .
                " --output $outdir/step2_regulons.csv --num_workers $num_workers";
    print "[INFO: $time] Now start run step2 may take a long time, please be patient. \n";
    print "CMD: $CMD2\n";
    system("$CMD2 >> $outdir/run_log.txt 2>\&1");
}

if($transpose){
    my $time = &format_datetime(localtime(time()));
    my $CMD3 = "$pyscenic aucell --output $outdir/step3_auc_mtx.csv --num_workers $num_workers --seed 100 --transpose $mtx $outdir/step2_regulons.csv";
    print "[INFO: $time] Now start run step3 may take a long time, please be patient. \n";
    print "CMD: $CMD3\n";
    system("$CMD3 >> $outdir/run_log.txt 2>\&1");
}else{
    my $time = &format_datetime(localtime(time()));
    my $CMD3 = "$pyscenic aucell --output $outdir/step3_auc_mtx.csv --num_workers $num_workers --seed 100 $mtx $outdir/step2_regulons.csv";
    print "[INFO: $time] Now start run step3 may take a long time, please be patient. \n";
    print "CMD: $CMD3\n";
    system("$CMD3 >> $outdir/run_log.txt 2>\&1");
}

# 经过 step3 后, 出来的 step3_auc_mtx.csv 文件是以逗号分隔, 行名为 Regulon(+) 的形式, 将它处理成制表符分隔, 行名处理成 Regulon(159g) 的形式
my $time = &format_datetime(localtime(time()));
print "[INFO: $time] All steps of pyscenic is done, Now start treat all outpt from pyscenic. \n";
my %hash;

open IN,"$outdir/step2_regulons.csv";
while(<IN>){
    chomp;
    next if $. <= 3;
    my($TF) = (split/,/,$_)[0];
    my($context) = $_ =~ m/"\[(.*)\]"/;
    my @gene_list = $context =~ /\'([^\s]*)\'/g;
    push @{$hash{$TF}},@gene_list;
}
close IN;

my %transname;
for my $key(keys %hash){
    my %uniq;
    my $uniq_gene_num = grep{++$uniq{$_} < 2} @{$hash{$key}};
    $transname{$key} = "$key". "(" . "$uniq_gene_num" ."g)";
}

open IN,"$outdir/step3_auc_mtx.csv";
open OUT,">$outdir/finalout_auc_score_mtx.xls";
while(<IN>){
    chomp;
    my $line;
    if($. == 1){
        s/\./\-/g if $_ =~ /\./;
        s/\,/\t/g if $_ =~ /\,/;
        $line = $_;
    }else{
        my($row_name, $expr) = split/,/,$_,2;
        (my $gene = $row_name) =~ s/\(\+\)//;
        $expr =~ s/\,/\t/g if $expr =~ /\,/;
        my $new_rowname;
        if(exists $transname{$gene}){
            $new_rowname = $transname{$gene};
        }else{
            $new_rowname = $row_name;
        }
        
        $line = join "\t",($new_rowname,$expr);
    }
    print OUT "$line\n";
}
close IN;
close OUT;